#!/usr/bin/env python3

import os
import mmap
import struct
import time
import numpy as np
import cv2
import rospy
import threading
from sensor_msgs.msg import Image
from cv_bridge import CvBridge

class ImageAcquisition:
    def __init__(self):
        self.bridge = CvBridge()
        self.color_img = None
        self.depth_img = None
        self.lock = threading.Lock()
        
        # ROS订阅
        rospy.Subscriber("/camera/color/image_raw", Image, self.color_callback)
        rospy.Subscriber("/camera/depth/image_raw", Image, self.depth_callback)
        rospy.loginfo("ImageAcquisition initialized, waiting for images...")
        
    def color_callback(self, msg):
        try:
            with self.lock:
                self.color_img = msg
                self.try_process_images()
        except Exception as e:
            rospy.logerr(f"Color error: {e}")

    def depth_callback(self, msg):
        try:
            with self.lock:
                self.depth_img = msg
                self.try_process_images()
        except Exception as e:
            rospy.logerr(f"Depth error: {e}")

    def try_process_images(self):
        if self.color_img is not None and self.depth_img is not None:
            success = self.shm_manager.write_images(self.color_img, self.depth_img)
            if success:
                rospy.loginfo_throttle(5, "Image pair written to shared memory")
            
            # 重置
            self.color_img = None
            self.depth_img = None

    def shutdown(self):
        """关闭时清理"""
        self.shm_manager.close()
        rospy.loginfo("ImageAcquisition shutdown")

if __name__ == "__main__":
    rospy.init_node("image_acquisition")
    acquisition = ImageAcquisition()
    rospy.on_shutdown(acquisition.shutdown)
    rospy.spin()